from p_tuning.train import TrainBertPTun

from pet.train import TrainBertPet


def main(use_type="pet"):
    if use_type == "pet":
        train_model = TrainBertPet()
    else:
        train_model = TrainBertPTun()
    train_model.running()


if __name__ == '__main__':
    # use_type = "pet"
    use_type = "p_tuning"
    main(use_type)
